from typing import List, Dict, Any, Optional
import logging
from neo4j import GraphDatabase
from py2neo import Graph
import numpy as np
from .base import BaseGraphStore, Node, Edge, SubGraph


class Neo4jGraphStore(BaseGraphStore):
    """Neo4j图数据库实现"""
    
    def __init__(self):
        self.driver = None
        self.graph = None
        self.logger = logging.getLogger(__name__)
        
    def initialize(self, config: Dict[str, Any]) -> None:
        """初始化Neo4j连接"""
        try:
            uri = config["uri"]
            username = config["username"]
            password = config["password"]
            
            # 初始化Neo4j驱动
            self.driver = GraphDatabase.driver(uri, auth=(username, password))
            
            # 初始化py2neo Graph对象
            self.graph = Graph(uri, auth=(username, password))
            
            self.logger.info("Neo4j连接初始化成功")
            
        except Exception as e:
            self.logger.error(f"Neo4j连接初始化失败: {str(e)}")
            raise
            
    def add_node(self, node: Node) -> None:
        """添加节点到图数据库"""
        try:
            # 构建Cypher查询
            properties = dict(node.properties)
            if node.embedding is not None:
                properties["embedding"] = node.embedding
                
            query = (
                f"CREATE (n:{node.type} {{id: $id}}) "
                "SET n += $properties"
            )
            
            # 执行查询
            with self.driver.session() as session:
                session.run(
                    query,
                    id=node.id,
                    properties=properties
                )
                
        except Exception as e:
            self.logger.error(f"添加节点失败: {str(e)}")
            raise
            
    def add_edge(self, edge: Edge) -> None:
        """添加边到图数据库"""
        try:
            # 构建Cypher查询
            query = (
                "MATCH (source {id: $source_id}), (target {id: $target_id}) "
                f"CREATE (source)-[r:{edge.type} $properties]->(target)"
            )
            
            # 执行查询
            with self.driver.session() as session:
                session.run(
                    query,
                    source_id=edge.source_id,
                    target_id=edge.target_id,
                    properties=edge.properties
                )
                
        except Exception as e:
            self.logger.error(f"添加边失败: {str(e)}")
            raise
            
    def get_node(self, node_id: str) -> Optional[Node]:
        """获取指定ID的节点"""
        try:
            # 构建Cypher查询
            query = (
                "MATCH (n {id: $id}) "
                "RETURN n"
            )
            
            # 执行查询
            with self.driver.session() as session:
                result = session.run(query, id=node_id)
                record = result.single()
                
                if record is None:
                    return None
                    
                # 转换为Node对象
                node = record["n"]
                properties = dict(node)
                embedding = properties.pop("embedding", None)
                
                return Node(
                    id=node_id,
                    type=next(iter(node.labels)),
                    properties=properties,
                    embedding=embedding
                )
                
        except Exception as e:
            self.logger.error(f"获取节点失败: {str(e)}")
            raise
            
    def get_neighbors(self, node_id: str, edge_types: Optional[List[str]] = None) -> List[Node]:
        """获取节点的邻居节点"""
        try:
            # 构建Cypher查询
            if edge_types:
                edge_pattern = "|".join(f":{t}" for t in edge_types)
                query = (
                    "MATCH (n {id: $id})-[r" + edge_pattern + "]-(neighbor) "
                    "RETURN DISTINCT neighbor"
                )
            else:
                query = (
                    "MATCH (n {id: $id})-[]-(neighbor) "
                    "RETURN DISTINCT neighbor"
                )
            
            # 执行查询
            neighbors = []
            with self.driver.session() as session:
                result = session.run(query, id=node_id)
                
                for record in result:
                    node = record["neighbor"]
                    properties = dict(node)
                    embedding = properties.pop("embedding", None)
                    
                    neighbors.append(Node(
                        id=properties.pop("id"),
                        type=next(iter(node.labels)),
                        properties=properties,
                        embedding=embedding
                    ))
                    
            return neighbors
            
        except Exception as e:
            self.logger.error(f"获取邻居节点失败: {str(e)}")
            raise
            
    def search_nodes(self, query: str, node_types: Optional[List[str]] = None) -> List[Node]:
        """搜索节点"""
        try:
            # 构建Cypher查询
            if node_types:
                type_pattern = "|".join(f":{t}" for t in node_types)
                cypher_query = (
                    f"MATCH (n{type_pattern}) "
                    "WHERE n.content CONTAINS $query "
                    "RETURN n"
                )
            else:
                cypher_query = (
                    "MATCH (n) "
                    "WHERE n.content CONTAINS $query "
                    "RETURN n"
                )
            
            # 执行查询
            nodes = []
            with self.driver.session() as session:
                result = session.run(cypher_query, query=query)
                
                for record in result:
                    node = record["n"]
                    properties = dict(node)
                    embedding = properties.pop("embedding", None)
                    
                    nodes.append(Node(
                        id=properties.pop("id"),
                        type=next(iter(node.labels)),
                        properties=properties,
                        embedding=embedding
                    ))
                    
            return nodes
            
        except Exception as e:
            self.logger.error(f"搜索节点失败: {str(e)}")
            raise
            
    def get_subgraph(self, node_ids: List[str], depth: int = 1) -> SubGraph:
        """获取以给定节点为中心的子图"""
        try:
            # 构建Cypher查询获取节点和边
            query = (
                "MATCH path = (n)-[*0.." + str(depth) + "]-(m) "
                "WHERE n.id IN $node_ids "
                "RETURN path"
            )
            
            nodes = set()
            edges = set()
            
            # 执行查询
            with self.driver.session() as session:
                result = session.run(query, node_ids=node_ids)
                
                for record in result:
                    path = record["path"]
                    
                    # 收集路径中的节点
                    for node in path.nodes:
                        properties = dict(node)
                        embedding = properties.pop("embedding", None)
                        
                        nodes.add(Node(
                            id=properties.pop("id"),
                            type=next(iter(node.labels)),
                            properties=properties,
                            embedding=embedding
                        ))
                    
                    # 收集路径中的边
                    for rel in path.relationships:
                        edges.add(Edge(
                            source_id=rel.start_node["id"],
                            target_id=rel.end_node["id"],
                            type=rel.type,
                            properties=dict(rel)
                        ))
            
            return SubGraph(
                nodes=list(nodes),
                edges=list(edges)
            )
            
        except Exception as e:
            self.logger.error(f"获取子图失败: {str(e)}")
            raise
            
    def __del__(self):
        """清理资源"""
        if self.driver:
            self.driver.close() 